import json
import os
import argparse

if __name__ == '__main__':
    test_json_raw_f = json.load(open("data/annotations/front_test.json", "r"))
    test_json_f = json.load(open("/root/mmdetection/results/" + "front_ensemble.json", "r"))

    test_json_raw_s = json.load(open("data/annotations/side_test.json", "r"))
    test_json_s = json.load(open("/root/mmdetection/results/" + "side_ensemble.json", "r"))
    submit_file_name = "submit.csv"
    submit_path = 'submit/'
    os.makedirs(submit_path, exist_ok=True)

    csv_file = open(submit_path + submit_file_name, 'w')
    csv_file.write("name,image_id,confidence,xmin,ymin,xmax,ymax\n")

    imgid2anno_f = {}
    imgid2name_f = {}
    imgid2size_f = {}
    for imageinfo in test_json_raw_f['images']:
        imgid = imageinfo['id']
        imgid2name_f[imgid] = imageinfo['file_name']
        imgid2size_f[imgid] = (imageinfo['height'], imageinfo['width'])
    for anno in test_json_f:
        img_id = anno['image_id']
        if img_id not in imgid2anno_f:
            imgid2anno_f[img_id] = []
        imgid2anno_f[img_id].append(anno)
    for imgid, annos in imgid2anno_f.items():
        for anno in annos:
            confidence = anno['score']
            xmin, ymin, w, h = anno['bbox']
            xmax = xmin + w
            ymax = ymin + h
            xmin = int(xmin + 0.5)
            ymin = int(ymin + 0.5)
            xmax = min(int(xmax + 0.5), imgid2size_f[imgid][1])
            ymax = min(int(ymax + 0.5), imgid2size_f[imgid][0])

            class_name = 'target'
            image_name = imgid2name_f[imgid]
            image_id = image_name.split('.')[0] + '.xml'
            csv_file.write(
                class_name + ',' + image_id + ',' + str(confidence) + ',' + str(xmin) + ',' + str(ymin) + ',' + str(
                    xmax) + ',' + str(ymax) + '\n')

    imgid2anno_s = {}
    imgid2name_s = {}
    imgid2size_s = {}
    for imageinfo in test_json_raw_s['images']:
        imgid = imageinfo['id']
        imgid2name_s[imgid] = imageinfo['file_name']
        imgid2size_s[imgid] = (imageinfo['height'], imageinfo['width'])
    for anno in test_json_s:
        img_id = anno['image_id']
        if img_id not in imgid2anno_s:
            imgid2anno_s[img_id] = []
        imgid2anno_s[img_id].append(anno)
    for imgid, annos in imgid2anno_s.items():
        for anno in annos:
            confidence = anno['score']
            xmin, ymin, w, h = anno['bbox']
            xmax = xmin + w
            ymax = ymin + h
            xmin = int(xmin + 0.5)
            ymin = int(ymin + 0.5)
            xmax = min(int(xmax + 0.5), imgid2size_s[imgid][1])
            ymax = min(int(ymax + 0.5), imgid2size_s[imgid][0])

            class_name = 'target'
            image_name = imgid2name_s[imgid]
            image_id = image_name.split('.')[0] + '.xml'
            csv_file.write(
                class_name + ',' + image_id + ',' + str(confidence) + ',' + str(xmin) + ',' + str(ymin) + ',' + str(
                    xmax) + ',' + str(ymax) + '\n')
    csv_file.close()
